import csv
import os
import time


# 控制类
class Controller:
    def __init__(self, count, packageName):
        self.counter = count
        self.alldata = [('timestamp', 'traffic')]
        self.packageName = packageName
        self.receive = ''
        self.receive2 = ''
        self.transmit = ''
        self.transmit2 = ''

    #字符串转换成数字
    def atoi(self, num_string):
        if num_string == '':
            return 0
        else:
            try:
                f = float(num_string)
                i = int(f)
            except:
                return 0
            else:
                return i

    # 单次测试过程
    def testprocess(self):
        # 执行获取进程的命令
        result = os.popen('adb shell ps | grep %s' % self.packageName)
        # 获取进程id
        pid = result.readlines()[0].split(" ")[5]
        # 获取进程ID使用的流量
        traffic = os.popen("adb shell cat /proc/" + pid + "/net/dev")
        for line in traffic.readlines():
            if 'eth0' in line:
                # 将所有空行换成#
                line = "#".join(line.split())
                # 按#号拆分，获取收到和发出的流量
                self.receive = line.split('#')[1]
                self.transmit = line.split('#')[9]
            elif 'eth1' in line:
                # 将所有空行换成#
                line = "#".join(line.split())
                # 按#号拆分，获取收到和发出的流量
                self.receive2 = line.split('#')[1]
                self.transmit2 = line.split('#')[9]
        # 计算所有流量之和
        print('receive---->', self.receive)
        print('transmit--->', self.transmit)
        print('self.receive2 --->>', self.receive2)
        print('self.transmit2', self.transmit2)
        alltraffic = self.atoi(self.receive) + self.atoi(self.transmit) + self.atoi(self.receive2) + self.atoi(
            self.transmit2)
        print('alltraffic', alltraffic)
        # 按KB计算流量值
        alltraffic = alltraffic / 1024
        currenttime = self.getCurrentTime()
        # 将获取到的数据存到数组中
        self.alldata.append((currenttime, int(alltraffic)))

    # 多次执行过程
    def run(self):
        while self.counter > 0:
            self.testprocess()
            self.counter -= 1
            # 每5秒采集一次数据
            time.sleep(5)

    def getCurrentTime(self):
        currentTime = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
        return currentTime

    def SaveDataToCSV(self):
        with open('traffic.csv', 'w+') as fp:
            writer = csv.writer(fp)
            writer.writerows(self.alldata)


if __name__ == '__main__':
    controller = Controller(count=5, packageName='com.android.browser')
    controller.run()
    controller.SaveDataToCSV()
